package com.lun.easy;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.stream.Collectors;

import com.lun.util.BinaryTree.TreeNode;

public class AverageOfLevelsInBinaryTree {

	// 方法一：BFS
	public List<Double> averageOfLevels(TreeNode root) {
		LinkedList<TreeNode> queue = new LinkedList<>();

		queue.offer(root);
		List<Double> result = new ArrayList<>();

		while (!queue.isEmpty()) {

			int size = queue.size(), originSize = size;
			long sum = 0;

			for (int i = 0; i < size; i++) {
				TreeNode node = queue.poll();

				sum += node.val;

				if (node.left != null)
					queue.offer(node.left);

				if (node.right != null)
					queue.offer(node.right);

			}
			result.add(sum * 1.0 / originSize);
		}

		return result;
	}

	// 方法二：DFS
	public List<Double> averageOfLevels2(TreeNode root) {
		List<long[]> levelSumAndCounts = new ArrayList<>();
		averageOfLevels2(root, 0, levelSumAndCounts);
		return levelSumAndCounts.stream().map(a -> a[0] * 1.0 / a[1]).collect(Collectors.toList());
	}

	private void averageOfLevels2(TreeNode node, int level, List<long[]> levelSumAndCounts) {
		if (node == null)
			return;

		if (levelSumAndCounts.size() == level) {
			levelSumAndCounts.add(new long[] {node.val, 1});
		} else {
			long[] temp = levelSumAndCounts.get(level);
			temp[0] += node.val;
			temp[1]++;
		}

		averageOfLevels2(node.left, level + 1, levelSumAndCounts);
		averageOfLevels2(node.right, level + 1, levelSumAndCounts);

	}

}
